from get_codes import *
from get_price import Get_price
import sys 
sys.path.append(r'code/data_capture')
from data_contrate import Data_consentrate
from renew_price import Renew_price
import akshare as ak
import numpy as np


def get_data():
    ### 登陆系统 ####
    bs.login()

    ###获取股票代码
    GC = Get_codes()
    GC.get_sz50_codes()
    GC.get_hs300_codes()
    GC.get_zz500_codes()
    GC.get_zz1000_codes()

    sz50 = pd.read_csv(r'data/stock_data/code_date/sz50_stocks.csv')
    hs300 = pd.read_csv(r'data/stock_data/code_date/hs300_stocks.csv')
    zz500 = pd.read_csv(r'data/stock_data/code_date/zz500_stocks.csv')
    zz1000 = pd.read_csv(r'data/stock_data/code_date/zz1000_stocks.csv')

    sz50 = sz50['code'].values
    hs300 = hs300['code'].values
    zz500 = zz500['code'].values
    zz1000 = zz1000['code'].values
    
    # G = Get_price(path=r'data/stock_data/min_price', start_date='2001-01-01', freq='5')
    G = Get_price()
    G.get_price(np.concatenate([sz50,hs300,zz500,zz1000 ], axis=0))
    
    bs.logout()


def renew_data():
    data_file = os.listdir(r'data/stock_data/daily_price')
    lst = [x[:-9] for x in data_file]
    bs.login()
    R = Renew_price()
    R.get_price(lst)
    bs.logout()

def catdata():
    lst = ['open', 'high', 'low', 'close', 'volume', 'vol', 'lagret']
    for i in lst:
        D = Data_consentrate()#root=r'data/stock_data/min_price')
        D.consentrate(key=i)
        
def clean_data(root=None):
    sz50 = pd.read_csv(r'data/stock_data/code_date/sz50_stocks.csv')
    hs300 = pd.read_csv(r'data/stock_data/code_date/hs300_stocks.csv')
    zz500 = pd.read_csv(r'data/stock_data/code_date/zz500_stocks.csv')
    zz1000 = pd.read_csv(r'data/stock_data/code_date/zz1000_stocks.csv')

    sz50 = sz50['code'].values
    hs300 = hs300['code'].values
    zz500 = zz500['code'].values
    zz1000 = zz1000['code'].values
    if root is None:
        root = 'data/stock_data/daily_price'
    data_file = os.listdir(root)
    data_lst = [x[:-9] for x in data_file]
    unino = []
    code_lst = [sz50, hs300,zz500,zz1000]
    for i in code_lst:
        unino = np.union1d(unino, i)
    drop_lst = np.setdiff1d(data_lst, unino)
    try:
        hold_lst = pd.read_csv('result/real/10/pos/hold_pos.csv', index_col=0)
        hold_lst = hold_lst.iloc[-1,:]
        hold_lst = hold_lst[hold_lst==0]
        hold_lst = hold_lst.index
    except Exception:
        hold_lst = []
    for i in drop_lst:
        if i in hold_lst:
            os.remove(os.path.join(root,f'{i}.pkl.gzip'))


# get_data()
renew_data()
catdata()
clean_data()